# -*- coding:utf-8 -*-

import re

import jieba.posseg

jieba.load_userdict("data/person.txt")

STOP_WORDS = set([w.strip() for w in open("data/stopwords.txt", encoding='utf-8').readlines()])


class MyChapters(object):
    def __init__(self, chapter_list):
        self.chapter_list = chapter_list

    def __iter__(self):
        for chapter in self.chapter_list:
            yield cut_words_with_pos(chapter)


def split_by_chapter(filepath):
    """
    将文档按章节切割
    chapter_list： [[第一章内容],[第二章内容],[..],[..],...]
    """
    with open(filepath, encoding='utf-8') as f:
        text = f.read()
        chapter_list = re.split(r'第.{1,3}章\n', text)[1:]
    return chapter_list


def cut_words_with_pos(text):
    """
    按照词性对每一章节进行切词
    :param text:
    :return:
    """
    seg = jieba.posseg.cut(text)
    res = []
    for i in seg:
        if i.flag in ["a", "v", "x", "n", "an", "vn", "nz", "nt", "nr"] and is_fine_word(i.word):
            res.append(i.word)

    return res


def is_fine_word(word, min_length=2):
    """
    过滤词长，过滤停用词，只保留中文 
    :param word: 
    :param min_length: 
    :return: 
    """
    rule = re.compile(r"^[\u4e00-\u9fa5]+$")
    if len(word) >= min_length and word not in STOP_WORDS and re.search(rule, word):
        return True
    else:
        return False
